from casevo import ModelBase
from mesa.space import NetworkGrid
from mesa.time import RandomActivation
import networkx as nx
import json
import os
import numpy as np
from casevo.agent_base import AgentBase
from casevo import TotLog
from .llm_interface import BaichuanLLM
from .agent import SimpleAgent

class SimpleABMModel(ModelBase):
    """简化版ABM模型"""
    
    def __init__(self, config_path="config.json"):
        """
        初始化模型
        
        Args:
            config_path: 配置文件路径
        """
        # 加载配置
        self.config = self._load_config(config_path)
        
        # 设置随机种子
        self.random.seed(self.config.get('random_seed', 42))
        
        # 创建网格
        graph = nx.watts_strogatz_graph(
            self.config['num_agents'],
            self.config['network']['k'],
            self.config['network']['p']
        )
        
        # 初始化父类
        super().__init__(graph)
        
        # 加载代理配置文件
        self.profiles = self._load_profiles()
        
        # 创建代理
        self.schedule = RandomActivation(self)
        for i in range(self.config['num_agents']):
            # 使用循环方式获取配置文件
            profile = self.profiles[i % len(self.profiles)] if self.profiles else None
            agent = SimpleAgent(i, self, profile)
            self.schedule.add(agent)
            
        # 将代理放置到网格中
        for i, agent in enumerate(self.schedule.agents):
            self.grid.place_agent(agent, i)
            
    def _load_config(self, config_path):
        """加载配置文件"""
        with open(config_path, 'r', encoding='utf-8') as f:
            return json.load(f)
            
    def _load_profiles(self):
        """加载代理配置文件"""
        profiles_path = self.config.get('agent', {}).get('profiles_file', 'agent_profiles.json')
        if os.path.exists(profiles_path):
            with open(profiles_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        return []
        
    def step(self):
        """执行一步"""
        self.schedule.step()
        
    def get_agent_states(self):
        """获取所有代理的状态"""
        return [agent.get_state() for agent in self.schedule.agents]

    def _setup_logging(self):
        """设置日志记录"""
        # 创建日志目录
        if not os.path.exists('log'):
            os.makedirs('log')
            
        # 初始化模型日志
        TotLog.model_log = []
        
        # 初始化代理日志
        TotLog.agent_log = [[] for _ in range(self.config['model']['num_agents'])]
            
    def _generate_summary(self):
        """生成模型状态摘要"""
        opinions = [agent.opinion for agent in self.schedule.agents]
        confidences = [agent.confidence for agent in self.schedule.agents]
        
        return {
            'step': self.step_count,
            'total_agents': len(self.schedule.agents),
            'avg_opinion': sum(opinions) / len(opinions),
            'avg_confidence': sum(confidences) / len(confidences),
            'opinion_std': np.std(opinions),
            'confidence_std': np.std(confidences)
        }
            
    def run_model(self, n_steps=None):
        """运行模型"""
        if n_steps is None:
            n_steps = self.config['simulation']['max_steps']
            
        for _ in range(n_steps):
            self.step()
            
            # 定期记录日志
            if self.step_count % self.config['simulation']['log_interval'] == 0:
                self._save_log()
                
    def _save_log(self):
        """保存日志"""
        log_file = f"log/step_{self.step_count}.json"
        summary = self._generate_summary()
        
        with open(log_file, 'w', encoding='utf-8') as f:
            json.dump(summary, f, ensure_ascii=False, indent=2) 